#ifndef EPUBDUKPT_IMPL_DUKPT_HPP
#define EPUBDUKPT_IMPL_DUKPT_HPP

SJR_NAMESPACE_BEGIN

class EpubDukpt final {
public:
  constexpr static int KSI_BLOCK_SIZE = 5;
  constexpr static int BDK_BLOCK_SIZE = 16;
  constexpr static int DID_BLOCK_SIZE = 3;
  constexpr static int KSN_BLOCK_SIZE = 10;
  constexpr static int KEY_BLOCK_SIZE = 8;

  using IKSN     = unsigned char[KSN_BLOCK_SIZE];
  using IKCV     = unsigned char[KEY_BLOCK_SIZE];
  using IPEK     = unsigned char[BDK_BLOCK_SIZE];
  using K1       = unsigned char[KEY_BLOCK_SIZE];
  using K2       = unsigned char[KEY_BLOCK_SIZE];
  using K3       = unsigned char[BDK_BLOCK_SIZE];
  using DeviceID = unsigned char const[DID_BLOCK_SIZE];

  constexpr static unsigned char IV[] = {0xC0, 0xC0, 0xC0, 0xC0, 0x00, 0x00, 0x00, 0x00, 0xC0, 0xC0,
    0xC0, 0xC0, 0x00, 0x00, 0x00, 0x00, 0xC0, 0xC0, 0xC0, 0xC0, 0x00, 0x00, 0x00, 0x00};
  constexpr static unsigned char KCV[] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

  class DataSet final {
  public:
    friend EpubDukpt;

    DataSet() : mIKSN(), mIPEK(), mIKCV(), mKey1(), mKey2(), mKey3() {}

    auto iksnbytes() const -> int {
      return sizeof(mIKSN);
    }

    auto ipekbytes() const -> int {
      return sizeof(mIPEK);
    }

    auto ikcvbytes() const -> int {
      return sizeof(mIKCV);
    }

    auto iksn() const -> unsigned char const * {
      return mIKSN;
    }

    auto ipek() const -> unsigned char const * {
      return mIPEK;
    }

    auto ikcv() const -> unsigned char const * {
      return mIKCV;
    }

  private:
    IKSN mIKSN;
    IPEK mIPEK;
    IKCV mIKCV;

    K1 mKey1;
    K2 mKey2;
    K3 mKey3;
  };

  explicit EpubDukpt(EpubContext *context) : mContext(context) {}

  auto disperse(DeviceID did, DataSet *dataset) -> bool {
    auto &&params = (*mContext).params();
    auto &&ksi    = params.ksi();
    auto &&bdk    = params.bdk();

    if (ksi.size() != KSI_BLOCK_SIZE)
      return false;
    if (bdk.size() != BDK_BLOCK_SIZE)
      return false;
    if (!handleDukptZMK(mContext, &params))
      return false;
    if (!handleDukptKSN(&params, dataset, did))
      return false;
    if (!handleDukptCalcK1(mContext, &params, dataset))
      return false;
    if (!handleDukptCalcK2(mContext, &params, dataset))
      return false;
    if (!handleDukptCalcKcv(mContext, dataset))
      return false;

    auto &&step = 0;
    memcpy(&(*dataset).mIPEK[step], (*dataset).mKey1, sizeof((*dataset).mKey1));
    step += sizeof((*dataset).mKey1);
    memcpy(&(*dataset).mIPEK[step], (*dataset).mKey2, sizeof((*dataset).mKey2));
    return true;
  }

private:
  EpubContext *mContext;

  static auto handleDukptZMK(EpubContext *context, EpubContextParams *params) -> bool {
    if ((*params).zmk().empty())
      return true;

    auto &&ret = bool(true);
    auto &&bdk = std::vector<unsigned char>((*params).zmk().size());
    auto &&key = (*context).createSchedule((*params).zmk());

    for (int i = 0; i < bdk.size(); i += 8) {
      if ((*context).encrypt(key, &(*params).bdk()[i], &bdk[i]))
        continue;
      ret = false;
      break;
    }

    (*params).setBDK(bdk.data(), static_cast<int>(bdk.size()));
    (*context).destroySchedule(key);
    return ret;
  }

  static auto handleDukptKSN(EpubContextParams *params, DataSet *dataset, DeviceID did) -> bool {
    auto &&ksiLen = (*params).ksi().size();

    if ((did[2] & 0x01u) != 0)
      return false;
    for (int i = 0; i < ksiLen; ++i)
      (*dataset).mIKSN[i] = (*params).ksi()[i];

    (*dataset).mIKSN[5] = ((did[0] & 0x0Fu) << 4) + ((did[1] & 0xF0u) >> 4);
    (*dataset).mIKSN[6] = ((did[1] & 0x0Fu) << 4) + ((did[2] & 0xF0u) >> 4);
    (*dataset).mIKSN[7] = ((did[2] & 0x0Fu) << 4);
    return true;
  }

  static auto handleDukptCalcK1(EpubContext *context, EpubContextParams *params, DataSet *dataset)
    -> bool {
    auto &&ret = false;
    auto &&key = (*context).createSchedule((*params).bdk());

    do {
      if (!(*context).encrypt(key, (*dataset).mIKSN, (*dataset).mKey1))
        break;
      ret = true;
    } while (false);

    (*context).destroySchedule(key);
    return ret;
  }

  static auto handleDukptCalcK2(EpubContext *context, EpubContextParams *params, DataSet *dataset)
    -> bool {
    auto &&ret = false;

    for (int i = 0; i < (*params).bdk().size(); ++i)
      (*dataset).mKey3[i] = (*params).bdk()[i] ^ IV[i];

    auto &&key = (*context).createSchedule((*dataset).mKey3, sizeof((*dataset).mKey3));

    do {
      if (!(*context).encrypt(key, (*dataset).mIKSN, (*dataset).mKey2))
        break;
      ret = true;
    } while (false);

    (*context).destroySchedule(key);
    return ret;
  }

  static auto handleDukptCalcKcv(EpubContext *context, DataSet *dataSet) -> bool {
    auto &&ret = true;
    auto &&key = (*context).createSchedule((*dataSet).mIPEK, sizeof((*dataSet).mIPEK));

    for (int i = 0; i < sizeof(KCV); i += 8) {
      if ((*context).encrypt(key, &KCV[i], &(*dataSet).mIKCV[i]))
        continue;
      ret = false;
      break;
    }

    (*context).destroySchedule(key);
    return ret;
  }
};

SJR_NAMESPACE_END

#endif//EPUBDUKPT_IMPL_DUKPT_HPP
